#!/bin/bash

# UEFI installer
#  (c) copyright Michael Amadio, 2016, Gold Coast, Australia
#  ALL RIGHTS RESERVED
#  01micko@gmail.com
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#  
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#  
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
# 
# see /usr/share/doc/legal for full licence 
# first version
# wizard style script to install to UEFI bootable Puppy Linux to a USB
# flash drive
#-----------------------------------------------------------------------------#
# depends: 
# UEFI enabled kernel
# +gtkdialog >=0.8
# +gtkdialog-splash
# +Xdialog 
# +gparted
# +gettext
# +grubx*.efi.tar.xz from jamesbond
#-----------------------------------------------------------------------------#
ver=0.2

# vars
. /etc/DISTRO_SPECS
. /etc/rc.d/PUPSTATE
TMPMNT1=/mnt/UEFI1
mkdir -p "$TMPMNT1"
TMP=/tmp/UEFI$$
mkdir -p "$TMP"
ABTMSG=$(gettext "You aborted the installation.")
FULLINST=yes
MNTERR=$(gettext "Aborting... Failed to mount ")
CPERR=$(gettext "Copy error")
SYSDIR=`find /usr -type d -name 'syslinux' -maxdepth 4`
LABEL=""
DISTRO=""
VERSION=""

# functions====================================================================
error_func() {
	#rm -r "$TMP"
	retval="$3"
	color="$2"
	[ "$retval" ] || color=red retval=$2
	gtkdialog-splash -close box -bg $color -text "$1"
	[ "$retval" ] && exit "$retval" || exit 1
}

# check running state, frugal and live only supported
pupcheck() {
	case "$PUPMODE" in
	 2|3|6|7|77)error_func "Can't use this script, wrong pupmode" 255
	 ;;
	esac
}

# check Grub UEFI
grub2_chk() {
	GRUB_ERR="No suitable grub2 file found. Exiting"
	ls /usr/share/grub2-efi | grep -qE 'grubx[0-9][0-9].efi.tar.xz' || error_func "$GRUB_ERR" orange $?
}

# find usb drive
find_usb() {
	probedisk|grep 'usb' > $TMP/uefi-wiz-list
	if [ ! -s "$TMP/uefi-wiz-list" ];then
	  MSG2=\
$(gettext "A USB drive can not be found. Is it plugged in? Exiting now")
	  error_func "$MSG2" 3
	fi
}

# fix mbr
mbr_func () {
	device="$1"
	echo "======================"
	echo "mbr_func"
	echo "======================"
	dd if=${device} of=/tmp/mbr_old.bin bs=446 count=1
	
	dd if=${SYSDIR}/mbr.bin of=${device} bs=440 count=1 conv=notrunc
	retval="$?"
	return $retval
}

# gui 1
func_gui1() {
	echo "======================"
	echo "func_gui1"
	echo "======================"
	
	# construct gui 1
	DESC=$(gettext "USB device - ")
	while read line
	do dev="${line%%|*}"
	desc="${line##*|}"
	echo "<radiobutton>
	 <label>$DESC $dev $desc</label>
	 <variable>$dev</variable>
	</radiobutton>" >> $TMP/uefi-wiz-radbtn
	done <"$TMP"/uefi-wiz-list
	RADBTN=`cat $TMP/uefi-wiz-radbtn`
	export gui1='<window title="UEFI usb installer - 1 v-'$ver'" window-position="1">
	<vbox space-expand="true" space-fill="true">
	  <text use-markup="true"><label>"<b>'$(gettext "Welcome to the UEFI usb drive installer wizard.")'</b>"</label></text>
	  <hseparator></hseparator>
	  <text><label>'$(gettext "Please choose the drive you wish to install to below.")'</label></text>
	  <frame>
	   '$RADBTN'
	  </frame>
	  <hbox homogeneous="true"><button ok></button></hbox>
	</vbox>
	</window>'
	gtkdialog -p gui1 > "$TMP"/uefi-chosen
	CHOSEN=`grep "true" "$TMP"/uefi-chosen|cut -d '=' -f1`
	echo $CHOSEN
	if [ ! "$CHOSEN" ];then 
	error_func "$ABTMSG" yellow 0
	fi
	DESC2=`grep "$CHOSEN" "$TMP"/uefi-wiz-list|rev|cut -d '|' -f1|rev` 
	# splash some info
	DMSG1=$(gettext "You have chosen ")
	DMSG2=$(gettext "When you click OK gparted will open to allow you to partition your \
	drive. You must make a FAT32 type partition. Be sure the drive is not mounted and be aware that all data on \
	the drive will be destroyed. You may want to write this down.")
	DMSG3=$(gettext "Be sure to mark the partition with the boot flag in gparted")
	export guigp='<window title="UEFI usb installer v-'$ver'" window-position="1">
	<vbox space-expand="true" space-fill="true">
	  <frame>
		<text><label>"'"$DMSG1 $CHOSEN - $DESC2"'"</label></text>
		<text><label>"'"$DMSG2"'"</label></text>
		<text use-markup="true"><label>"<b>'"$DMSG3"'</b>"</label></text>
	  </frame>
	<hbox homogeneous="true"><button ok></button></hbox>
	</vbox>
	</window>'
	eval `gtkdialog -p guigp -c`
	[ "$EXIT" = "abort" ] && error_func "$ABTMSG" yellow 0
	# run gparted on chosen partition
	GPERR=$(gettext "Gparted exited with an error. Exiting.")
	gparted "$CHOSEN"
	[ "$?" = 0 ] || error_func "$GPERR" $?
	dev1=${CHOSEN}1
	sync
	# check fs type
	FSNOT=$(gettext " is the wrong filesystem for that partition!")
	FATx2=`guess_fstype $dev1`
	[ "$FATx2" = "vfat" ] || error_func "$FATx2 $FSNOT" 255
	
	# check size, get away with 10M but we want at least 20, rec 32
	mount|grep -q "$TMPMNT1" || mount -t vfat "$dev1" "$TMPMNT1" -o utf8
	CSIZE=`df -m "$dev1"|awk '{print $2}'|tail -n1`
	if [ "$CSIZE" -lt 20 ];then
	SIZERR=$(gettext "Your boot partition is too small. Try again")
	umount "$TMPMNT1"
	error_func "$SIZERR" 77
	fi
	umount "$TMPMNT1"
	
	# check if boot flag set
	NOTBOOT=$(gettext " device is not bootable! Try again.")
	disktype "${CHOSEN}"|grep -q 'bootable'
	[ "$?" -ne 0 ] && error_func "$CHOSEN $NOTBOOT" 78
	
	# make bootable
	MBRF=(gettext "Failed to make bootable.")
	mbr_func "$CHOSEN"
	[ "$?" -ne 0 ] && error_func "$MBRF" "$?"
	
	FPART1=`probepart|grep -iE 'fat'|cut -d '|' -f1`
	
	rm -rf $TMP/UEFI* 2>/dev/null
	# sanity check
	FST1=`guess_fstype "$dev1"`
	echo $FST1
	MSG_FS1_1=$(gettext "is of type")
	MSG_FS1_2=$(gettext "it needs to be vfat. Exiting now.")
	if ! [ "`echo $FST1|grep -i 'fat'`" ];then
	error_func "$dev1 $MSG_FS1_1 $FST1 ${MSG_FS1_2}" 4
	fi
	
	sync
	
	func_frugal "$dev1"
}

# find and copy pup files
copy_func(){
	echo "======================"
	echo "copy_func"
	echo "======================"
	# note, initrd and vmlinuz
	# puppy_*_*.sfs, z(a,y)drv_*_*.sfs to the vfat partition
	SRCF="$1" dev1="$2"
	PUPPY=`ls "$SRCF"|grep '^puppy_'|grep 'sfs$'`
	DISTRO=`echo "$PUPPY"|cut -d _ -f2`
	VERSION=`echo "$PUPPY"|cut -d _ -f3`
	VERSION=${VERSION%.*}
	LABEL="$DISTRO $VERSION"
	echo "$LABEL"
	mount|grep -q "$TMPMNT1" || mount -t vfat "$dev1" "$TMPMNT1" -o utf8
	[ "$?" = 0 ] || error_func "$MNTERR ${dev1}" $?
	
	# Only support root of drive for now.. 
	HOMEDIR="$TMPMNT1"
	# copy *sfs
	for s in `ls $SRCF|grep 'sfs$'`;do
	  cp -a "${SRCF}"/"$s" $HOMEDIR
	done
	sync
	# copy kernel and initramfs
	for v in vmlinuz initrd.gz; do
	cp -a "${SRCF}"/"$v" $HOMEDIR
	done
	touch $HOMEDIR/USBFLASH # maybe not needed with new init?
	sync
	PID=`cat /tmp/u_pid`
	kill -9 $PID
	
	mount|grep -q "${SRCF}" && umount ${SRCF}
	
	boot_func "$dev1" ${HOMEDIR}
}

# gui 2
func_frugal() {
	echo "======================"
	echo "func_frugal"
	echo "======================"
	dev1="$1"
	FMSG1=$(gettext "Now it's time to specify your files.")
	FMSG2=$(gettext "Please choose where the Puppy files are located. \
The Puppy <u>must</u> be UEFI compatible!!!")
	export gui2='<window title="UEFI usb installer - frugal v-'$ver'" window-position="1">
	  <vbox>
	    <text use-markup="true"><label>"<b>'$FMSG1'</b>"</label></text>
	    <hseparator></hseparator>
	    <text use-markup="true"><label>"<i>'$FMSG2'</i>"</label></text>
	    <frame>
	     <hbox homogeneous="true">
	       <radiobutton><label>cd</label><variable>rb0</variable></radiobutton>
	       <radiobutton><label>iso image</label><variable>rb1</variable></radiobutton>
	       <radiobutton><label>files</label><variable>rb2</variable></radiobutton>
	     </hbox>
	    </frame>
	    <hbox homogeneous="true"><button ok></button></hbox>
	  </vbox>
	</window>'
	gtkdialog -p gui2 > "$TMP"/uefi-gui2
	. "$TMP"/uefi-gui2
	# choice of where files are located
	CPMSG1="$(gettext "Copying all the files to")"
	CPMSG2="$(gettext ". Please wait...")"
	case "$EXIT" in 
	OK)if [ "$rb0" = "true" ];then
		 gtkdialog-splash -bg yellow -close box -text "$CPMSG1 $dev1 $CPMSG2" &
		 CDPID=$!
		 echo $CDPID > /tmp/u_pid
		 mkdir -p /mnt/isomnt #mountpoint
		 mount -t iso9660 /dev/sr0 /mnt/cdmnt -o loop,ro
		 if [ "$?" -ne 0 ];then
		   sleep 1
		   kill -9 "$CDPID"
		   mount|grep -q 'cdmnt' && umount /mnt/cdmnt 2>/dev/null
		   FAILCDMSG="$(gettext "Failed to mount the CD! Is it there? Exiting")"
		   error_func "$FAILCDMSG" 9
		 fi
		 copy_func /mnt/cdmnt "$dev1"
		 if [ "$?" = 255 ];then
		   kill -9 "$CDPID"
		   mount|grep -q "$TMPMNT1" && umount "$TMPMNT1"
		   mount|grep -q 'cdmnt' && umount /mnt/cdmnt
		   error_func "$ABTMSG" $?
		 fi
		 
	   elif [ "$rb1" = "true" ];then
	     ISO="`Xdialog -title "UEFI - find ISO image" \
	     -backtitle "$(gettext "Please find the Puppy ISO image\n\
that you want to use for your install.")" \
	     -stdout --no-buttons --no-cancel --fselect "*" 0 0`"
	     [ ! $? -eq 0 ] && error_func "Something went wrong" 254
	     if [ ! `echo $ISO|grep 'iso$'` ];then
	       ISOMSG=$(gettext "is not an ISO image! Exiting")
	       error_func "$ISO $ISOMSG" 253
	     fi
		 gtkdialog-splash -bg yellow -close box -text "$CPMSG1 $dev1 $CPMSG2" &
		 ISOPID=$!
		 echo $ISOPID > /tmp/u_pid
		 mkdir -p /mnt/isomnt #mountpoint
		 mount -t iso9660 "$ISO" /mnt/isomnt -o loop,ro
		 if [ "$?" -ne 0 ];then
		   sleep 1
		   kill -9 "$ISOPID"
		   mount|grep -q 'isomnt' && umount /mnt/isomnt 2>/dev/null
		   CPISOMSG="$(gettext "Failed to mount the ISO! Is it valid? Exiting")"
		   error_func "$CPISOMSG" 252
		 fi
		 copy_func /mnt/isomnt "$dev1"
		 if [ "$?" = 255 ];then
		   kill -9 "$ISOPID"
		   mount|grep -q "$TMPMNT1" && umount "$TMPMNT1"
		   mount|grep -q 'isomnt' && umount /mnt/cdmnt
		   error_func "$ABTMSG" $?
		 fi
		 
		 
	   else # extracted from PUI, BK
	     SRCFPATH="`Xdialog -title "UEFI - find files" \
	     -backtitle "$(gettext "Please find the latest Puppy files\nvmlinuz, initrd.gz and related \
sfs files,\nTHEN HIGHLIGHT ANY ONE OF THEM in the right pane \nand click the OK button")" \
	     -stdout --no-buttons --no-cancel --fselect "*" 0 0`"
	     [ ! $? -eq 0 ] && exit
	     SRCPATH="${SRCFPATH%/*}"
	     if [ ! -f "${SRCPATH}"/initrd.gz ];then
	       FAIL0MSG="$(gettext "No Puppy files found in that directory! Exiting")"
	       error_func "$FAIL0MSG" 250
	     fi
	     echo "$SRCPATH"
	     gtkdialog-splash -bg yellow -close never -text "$CPMSG1 $dev1 $CPMSG2" &
		 CDPID=$!
		 echo $CDPID > /tmp/u_pid
		 copy_func $SRCPATH "$dev1"
	     if [ "$?" = 255 ];then
		   kill -9 "$CDPID"
		   mount|grep -q "$TMPMNT1" && umount "$TMPMNT1"
		   error_func "$ABTMSG" $?
		 elif [ "$?" -gt 0 ];then
		   sleep 1
		   kill -9 "$CDPID"
		   FAIL1MSG="$(gettext "Failed to copy the files! Is the path valid? Exiting")"
		   error_func "$FAIL1MSG" 251
		   exit
		 fi
		 kill -9 "$CDPID"
	   fi
	;;
	*)exit
	;;
	esac
}

# make syslinux bootable
sys_func () {
		echo "======================"
	echo "sys_func"
	echo "======================"
	dev1="$1" pic="$2"
	SYSPLASH=$(gettext "Now installing syslinux in to ")
	gtkdialog-splash -close never -timeout 3 -bg green -text "$SYSPLASH $dev1" &
	extlinux -i "$TMPMNT1"
	for b in vesamenu.c32 chain.c32;do 
	   cp ${SYSDIR}/$b "$TMPMNT1" 2>/dev/null #top level
	done
	sync

# build syslinux.conf	
cat > "$TMPMNT1"/syslinux.cfg <<SYS
#display help/boot.msg
default $DISTRO_FILE_PREFIX
prompt 1
timeout 100


ui vesamenu.c32
menu resolution 800 600
menu title $DISTRO_FILE_PREFIX Live
menu background ${pic}.png
menu tabmsg Press Tab to edit entry, Esc for boot prompt
menu color border 37;40  #80ffffff #00000000 std
menu color sel 7;37;40 #80ffffff #20ff8000 all
menu margin 1
menu rows 20
menu tabmsgrow 26
menu cmdlinerow -2
menu passwordrow 19
menu timeoutrow 28
menu helpmsgrow 30



label ${DISTRO_FILE_PREFIX}
linux vmlinuz
initrd initrd.gz
append pmedia=usbflash
menu label ${DISTRO_FILE_PREFIX}
text help
Start ${DISTRO_FILE_PREFIX} normally.
endtext


label ${DISTRO_FILE_PREFIX}-fsck
linux vmlinuz
initrd initrd.gz
append pfix=fsck pmedia=usbflash
menu label ${DISTRO_FILE_PREFIX} filesystem check
text help
Start ${DISTRO_FILE_PREFIX} normally with save filesystem check.
endtext


label ${DISTRO_FILE_PREFIX}-ram
linux vmlinuz
initrd initrd.gz
append pfix=ram pmedia=usbflash
menu label $DISTRO_FILE_PREFIX with no savefile
text help
Start ${DISTRO_FILE_PREFIX} with no savefile RAM only.
endtext


label ${DISTRO_FILE_PREFIX}-nox
linux vmlinuz
initrd initrd.gz
append pfix=nox pmedia=usbflash
menu label ${DISTRO_FILE_PREFIX} without graphical desktop
text help
Start ${DISTRO_FILE_PREFIX} in command-line mode (Linux console). 
Graphical desktop later can be started by typing "xwin".
endtext


menu separator

label ${DISTRO_FILE_PREFIX}-nokms
linux vmlinuz
initrd initrd.gz
append pfix=ram,nox pmedia=usbflash
menu label For machines with severe video problems
text help
Start ${DISTRO_FILE_PREFIX} without savefile, without KMS, and run xorgwizard 
to choose video resolutions before starting graphical desktop.
endtext
SYS
	
}

# grub2 - make bootable
boot_func() {
	echo "======================"
	echo "boot_func"
	echo "======================"
	dev1="$1"
	HOMEDIR="$2"
	GSPLASH=$(gettext "Now installing grub2 in to ")
	gtkdialog-splash -close never -timeout 3 -bg green -text "$GSPLASH $dev1" &
	GRUBDIR="/usr/share/grub2-efi"
	GRUB2="$GRUBDIR/grubx64.efi.tar.xz"
	mkdir -p ${HOMEDIR}/EFI/boot
	tar -xJvf $GRUB2 -C ${HOMEDIR}/EFI/boot/
	mv ${HOMEDIR}/EFI/boot/* ${HOMEDIR}/EFI/boot/bootx64.efi
	sync
	
# add boot image (png)
	pic=puppy
	case ${DISTRO_FILE_PREFIX} in
		[Tt]ahr*)pic=tahr;;
		[Ss]lacko*)pic=slacko;;
		[Xx]enial*)pic=xenial;;
	esac
	cp -a ${GRUBDIR}/${pic}.png ${HOMEDIR}/ 2>/dev/null

# build grub.cfg	
cat > ${HOMEDIR}/grub.cfg <<GRUB
insmod png
background_image /${pic}.png
set timeout=10
menuentry "Start $DISTRO_FILE_PREFIX" {
    linux /vmlinuz pmedia=usbflash
    initrd /initrd.gz
}
menuentry "Start $DISTRO_FILE_PREFIX - RAM only" {
    linux /vmlinuz pfix=ram pmedia=usbflash
    initrd /initrd.gz
}
menuentry "Start $DISTRO_FILE_PREFIX - No X" {
    linux /vmlinuz pfix=nox pmedia=usbflash
    initrd /initrd.gz
}
menuentry "Start $DISTRO_FILE_PREFIX - check filesystem" {
    linux /vmlinuz pfix=fsck pmedia=usbflash
    initrd /initrd.gz
}
menuentry "Start $DISTRO_FILE_PREFIX - No KMS" {
    linux /vmlinuz nomodeset pmedia=usbflash
    initrd /initrd.gz
}
menuentry "Shutdown" {
	halt
}
menuentry "Reboot" {
	reboot
}
GRUB

	# add syslinux if wanted
	ALTMSG=$(gettext "Would you like to add legacy BIOS boot capability to your installation?")
	Xdialog -title "Add Legacy BIOS" --yesno "$ALTMSG" 0 0 0
	case $? in
		0)sys_func $dev1 $pic;;
		*);;
	esac
	
	mount|grep -q "$TMPMNT1" && umount "$TMPMNT1"
	
	HAPPY=$(gettext "Everything looks good. Make sure $dev1 is unmounted and remove the media. \
After first boot you are given the option to make a \
'save file'. To keep your changes choose this or to boot to RAM every time don't cretate a save file. \
If you have Windows, be sure that you boot to Windows and disable hibernation, aka fast boot, aka hybdrid \
sleep, etc, and configure your UEFI to disable Secure Boot.")
	error_func "$HAPPY" green 0 #not an error, but a clean get out!
}


#main =========================================================================
pupcheck
grub2_chk
find_usb
func_gui1
[ -d "$TMP" ] && rm -r $TMP #cleanup
exit $?
